package org.wlgzs.sjl.红黑树;

import java.util.HashMap;

/**
 * ①创建RBTree，定义颜色
 * ②创建RBNode
 * ③辅助方法定义：parentOf(node)，isRed(node)，setRed(node)，setBlack(node)，inOrderPrint(RBNode tree)
 * ④左旋方法定义：leftRotate(node)
 * ⑤右旋方法定义：rightRotate(node)
 * ⑥公开插入接口方法定义：insert(K key, V value);
 * ⑦内部插入接口方法定义：insert(RBNode node);
 * ⑧修正插入导致红黑树失衡的方法定义：insertFIxUp(RBNode node);
 * ⑨测试红黑树正确性
 */
public class RBTree<K extends Comparable<K>, V> {

    //定义颜色常量
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    //红黑树的树根
    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    /**
     * 公开的插入接口
     *
     * @param key   键
     * @param value 值
     */
    public void insert(K key, V value) {
        RBNode node = new RBNode();
        node.setKey(key);
        node.setValue(value);
        node.setColor(RED);
        insert(node);
    }

    /**
     * 内部插入接口定义
     */
    private void insert(RBNode node) {
        //1.找到插入的位置
        RBNode parent = null;
        RBNode x = this.root;
        while (x != null) {
            parent = x;

            //a > b 则返回 1，否则返回 -1 ，相等返回0
            int cmp = node.key.compareTo(parent.key);

            if (cmp < 0) {
                x = x.left;
            } else if (cmp == 0) {
                parent.setValue(node.value);
                return;
            } else {
                x = x.right;
            }
        }

        node.parent = parent;

        if (parent != null) {
            if (node.key.compareTo(parent.key) < 0) {
                parent.left = node;
            } else {
                parent.right = node;
            }
        } else {
            this.root = node;
        }

        //插入之后需要进行修复红黑树，让红黑树再次平衡。
        insertFixUp(node);
    }

    /**
     * 插入后修复红黑树平衡的方法
     * |---情景1：红黑树为空树
     *              处理：直接把插入结点作为根结点。
     * |---情景2：插入节点的key已经存在
     *              处理：更新当前节点的值，为插入节点的值。
     * |---情景3：插入节点的父节点为黑色
     *              处理：直接插入即可。
     * |---情景4：插入节点的父节点为红色
     * |---情景4.1：叔叔节点存在，并且为红色（父-叔 双红）
     *                处理：父亲和叔叔节点改为黑色，爷爷节点改为红色，将爷爷节点设为当前节点继续平衡处理。
     * |---情景4.2：叔叔节点不存在，或者为黑色，父节点为爷爷节点的左子树
     * |---情景4.2.1：插入节点为其父节点的左子节点（LL情况）
     *                  处理：父亲节点改为黑色，爷爷节点改为红色，对爷爷节点进行右旋
     * |---情景4.2.2：插入节点为其父节点的右子节点（LR情况）
     *                  处理：对父亲节点进行左旋，得到LL情况，按照LL情况继续处理。
     * |---情景4.3：叔叔节点不存在，或者为黑色，父节点为爷爷节点的右子树
     * |---情景4.3.1：插入节点为其父节点的右子节点（RR情况）
     *                  处理：父亲节点改为黑色，爷爷节点改为红色，对爷爷节点进行左旋
     * |---情景4.3.2：插入节点为其父节点的左子节点（RL情况）
     *                  处理：对父亲节点进行右旋，得到RR情况，按照RR情况继续处理。
     */
    private void insertFixUp(RBNode node) {
        RBNode parent = parentOf(node);
        RBNode gparent = parentOf(parent);
        //情景4：存在父节点且父节点为红色
        if (parent != null && isRed(parent)) {
            //父节点是红色的，那么一定存在爷爷节点

            //父节点为爷爷节点的左子树
            if (parent == gparent.left) {
                RBNode uncle = gparent.right;
                //情景4.1：叔叔节点存在，并且为红色（父-叔 双红）
                //将父和叔染色为黑色，再将爷爷染红，并将爷爷设置为当前节点，进入下一次循环判断
                if (uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }

                //情景4.2：叔叔节点不存在，或者为黑色，父节点为爷爷节点的左子树
                if (uncle == null || isBlack(uncle)) {
                    //情景4.2.1：插入节点为其父节点的左子节点（LL情况）
                    //变色（父节点变黑，爷爷节点变红），右旋爷爷节点
                    if (node == parent.left) {
                        setBlack(parent);
                        setRed(gparent);
                        rightRotate(gparent);
                    }

                    //情景4.2.2：插入节点为其父节点的右子节点（LR情况）
                    //左旋（父节点）得到LL情况，当前节点设置为父节点，进入下一次循环
                    if (node == parent.right) {
                        leftRotate(parent);
                        insertFixUp(parent);
                        return;
                    }
                }
            } else {
                //父节点为爷爷节点的右子树
                RBNode uncle = gparent.left;
                //4.1：叔叔节点存在，并且为红色（父-叔 双红）
                //将父和叔染色为黑色，再将爷爷染红，并将爷爷设置为当前节点，进入下一次循环判断
                if (uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }

                //情景4.3：叔叔节点不存在，或者为黑色，父节点为爷爷节点的右子树
                if (uncle == null || isBlack(uncle)) {
                    //情景4.3.1：插入节点为其父节点的右子节点（RR情况）
                    //变色（父节点变黑，爷爷节点变红），左旋爷爷节点
                    if (node == parent.right) {
                        setBlack(parent);
                        setRed(gparent);
                        leftRotate(gparent);
                    }

                    //情景4.3.2：插入节点为其父节点的左子节点（RL情况）
                    //右旋（父节点）得到RR情况，当前节点设置为父节点，进入下一次循环
                    if (node == parent.left) {
                        rightRotate(parent);
                        insertFixUp(parent);
                        return;
                    }
                }
            }
        }

        setBlack(this.root);
    }

    /**
     * 公开的删除接口
     *
     * 1. 左右子树均不为空，进行删除转换;
     * 2. 删除转换后必定有一侧为空，如果该节点为红色节点，直接删除，并返回;
     * 3. 如果节点不为红色节点则分情况;
     *    3.1 其中之一为空，另一子树为红色节点（不存在另一节点为黑色节点非NIL的情况）,删除节点，子节点变色并顶替解决并返回；
     *    3.2 均为空；
     *
     * @param key   键
     */
    public boolean delete(K key) {
        RBNode node = findNode(key);

        if (node==null){
            return false;
        }

        //1.均不为空，则进行删除转换后继节点
        if(node.left != null && node.right != null){
            RBNode next = node.right;
            while(next.left != null){
                next = next.left;
            }
            node.key = next.key;
            node = next;
        }
        //2 节点为红色，直接删除
        if(node.color == RED){
            //node是红色节点且至少有一个子节点为NIL，所以不可能是根节点,且另一个子节点一定是NIL;
            if(node.parent.left == node){
                node.parent.left = null;
            } else {
                node.parent.right = null;
            }
            return true;
        }
        //3.1 node是黑色，其中之一为NIL，另一子树为红色节点：
        //左子树是红色节点
        if (node.left != null){
            node.left.color = BLACK;
            node.left.parent = node.parent;
            if (node == root){
                root = node.left;
            } else{
                if(node.parent.left == node){
                    node.parent.left = node.left;
                } else {
                    node.parent.right = node.left;
                }
            }
        } else if (node.right != null) {
            //右子树是红色节点
            node.right.color = BLACK;
            node.right.parent = node.parent;
            if (node == root){
                root = node.right;
            } else {
                if (node.parent.left == node){
                    node.parent.left = node.right;
                } else {
                    node.parent.right = node.right;
                }
            }
        }else{ //3.2 节点为黑色节点且子节点均为NIL
            deleteFixUp(node);
        }
        clearNode(node);
        return true;
    }

    /**
     * 删除红黑树自平衡方法
     *
     * |---情景1：左右子树均不为空
     *              处理：寻找后继节点进行替换操作，更新待平衡节点为替换后的节点。
     * |---情景2. 左右子树有一个为空
     * |---情景2.1：需要删除的节点为红色节点
     *                处理：直接删除。
     * |---情景2.2：需要删除的节点为黑色节点
     *                处理：如果一个子节点为红色，则把子节点变成黑色并替换该节点，删除替换后的节点。
     * |---情景3：两个子树均为null
     * |---情景3.1：需要删除的节点为红色节点
     *                处理：直接删除。
     * |---情景3.2：需要删除的节点为黑色节点
     * |---情景3.2.1：该节点的兄弟节点为红色
     *                  处理：将兄弟节点变为黑色，父亲节点变为红色，
     *                        若该节点为父节点的左子节点，则对父节点进行左旋处理，反之则进行右旋处理
     *                        此时情景转化为3.2.2。
     * |---情景3.2.2：该节点的兄弟节点为黑色
     * |---情景3.2.2.1：该节点的远侄子节点为红色且该节点为父节点的左子节点
     *                    处理：将父节点和兄弟节点的颜色对调。
     *                          对父节点进行左旋处理。
     *                          把远侄子节点变成黑色，并删除目标节点即可。
     * |---情景3.2.2.2：该节点的远侄子节点为红色且该节点为父节点的右子节点
     *                    处理：将父节点和兄弟节点的颜色对调。
     *                          对父节点进行右旋处理。
     *                          把远侄子节点变成黑色，并删除目标节点即可。
     * |---情景3.2.2.3：远侄子节点为黑色，近侄子节点为红色，该节点为父亲节点的左孩子
     *                    处理：将近侄子节点颜色设为父节点颜色。
     *                          将父节点颜色设为黑色。
     *                          将兄弟节点右旋。
     *                          将父节点节点左旋。
     * |---情景3.2.2.4：远侄子节点为黑色，近侄子节点为红色，该节点为父亲节点的右孩子
     *                    处理：将近侄子节点颜色设为父节点颜色。
     *                          将父节点颜色设为黑色。
     *                          将兄弟节点左旋。
     *                          将父节点节点右旋。
     * |---情景3.2.3：父亲节为红色，兄弟节点和兄弟节点的两个孩子（只能是NULL节点）都为黑色
     *                  处理：将父节点改成黑色，将兄弟节点改成红色。
     *                        然后删除目标节点即可。
     * |---情景3.2.4：父亲节点，兄弟节点和兄弟节点的两个孩子（只能为NULL节点）都为黑色
     *                  处理：将兄弟节点改成红色。
     *                        删除目标节点。
     *                        以父节点为目标节点继续进行自平衡操作
     *
     */
    private void deleteFixUp(RBNode node) {
        RBNode nodeTemp = node;
        //Fixup:
        while(nodeTemp != root) {
            boolean isLeft = nodeTemp.parent.left == nodeTemp;
            //定义兄弟节点
            RBNode sibling;
            if (isLeft) {
                sibling = nodeTemp.parent.right;
            } else {
                sibling = nodeTemp.parent.left;
            }
            //sibling is red ，转换成兄弟节点为黑色的情况；
            if (sibling.color == RED) {
                sibling.color = BLACK;
                nodeTemp.parent.color = RED;
                if (isLeft) {
                    leftRotate(nodeTemp.parent);
                    sibling = nodeTemp.parent.right;
                } else {
                    rightRotate(nodeTemp.parent);
                    sibling = nodeTemp.parent.left;
                }
            }
            //现在兄弟节点为黑色了；有以下几种情况：
            //1. 兄弟节点的内子节点为红色，通过旋转和变色达到平衡；
            //2. 兄弟节点的外子节点为红色，通过旋转和变色也能达到平衡；
            //3. 兄弟节点的内外子节点均不为红色，父节点为红色，通过变色可达到平衡；
            //4.  兄弟节点的内外子节点和父节点都不为红色，那么将兄弟节点子树的路径-1，树内部左右平衡，然后向上递归，寻找平衡机会；
            isLeft = nodeTemp.parent.left == nodeTemp;
            if (isLeft) {
                if (sibling.left != null && sibling.left.color == RED) {  //1
                    sibling.left.color = nodeTemp.parent.color;
                    nodeTemp.parent.color = BLACK;
                    rightRotate(sibling);
                    leftRotate(nodeTemp.parent);
                    break;
                } else if (sibling.right != null && sibling.right.color == RED) { // 2
                    sibling.color = nodeTemp.parent.color;
                    nodeTemp.parent.color = BLACK;
                    sibling.right.color = BLACK;
                    leftRotate(nodeTemp.parent);
                    break;
                } else if (nodeTemp.parent.color == RED) {// 3
                    sibling.color = RED;
                    nodeTemp.parent.color = BLACK;
                    break;
                } else { // 4
                    sibling.color = RED;
                    nodeTemp = nodeTemp.parent;
                }
            } else {
                if (sibling.right != null && sibling.right.color == RED){ //1
                    sibling.right.color = nodeTemp.parent.color;
                    nodeTemp.parent.color = BLACK;
                    leftRotate(sibling);
                    rightRotate(nodeTemp.parent);
                    break;
                } else if (sibling.left != null && sibling.left.color == RED){ // 2
                    sibling.color = nodeTemp.parent.color;
                    nodeTemp.parent.color = BLACK;
                    sibling.left.color = BLACK;
                    rightRotate(nodeTemp.parent);
                    break;
                } else if (nodeTemp.parent.color == RED) { // 3
                    sibling.color = RED;
                    nodeTemp.parent.color = BLACK;
                    break;
                } else { // 4
                    sibling.color = RED;
                    nodeTemp = nodeTemp.parent;
                }
            }
        }
        //remove
        if (node == root){
            root = null;
        } else {
            if(node.parent.left == node){
                node.parent.left = null;
            } else {
                node.parent.right = null;
            }
        }
    }

    public RBNode findNode(K key) {
        if(key == null){
            return null;
        }
        RBNode node =root;
        while(node != null){
            if(node.key.compareTo(key) < 0){
                node = node.right;
            } else if (node.key.compareTo(key) == 0){
                return node;
            } else {
                node = node.left;
            }
        }
        return null;
    }

    private void clearNode(RBNode node){
        node.left = null;
        node.right = null;
        node.parent = null;
        node.key = null;
        node = null;
    }

    /**
     * 左旋方法
     * 左旋示意图：左旋x节点
     * p                   p
     * |                   |
     * x                   y
     * / \         ---->   / \
     * lx  y               x   ry
     * / \             / \
     * ly  ry          lx  ly
     * <p>
     * 左旋做了几件事？
     * 1.将y的左子节点赋值给x的右边，并且把x设置为y的左子节点的父节点
     * 2.将x的父节点（非空时）指向y，更新y的父节点为x的父节点
     * 3.将y的左子节点指向x，更新x的父节点为y
     */
    private void leftRotate(RBNode x) {
        RBNode y = x.right;
        //将y的左子节点赋值给x的右边
        x.right = y.left;
        //并且把x设置为y的左子节点的父节点
        if (y.left != null) {
            y.left.parent = x;
        }

        //将x的父节点（非空时）指向y
        if (x.parent != null) {
            //如果x是parent左子树，则把y安放到parent的左边
            if (x.parent.left == x) {
                x.parent.left = y;
            } else {
                x.parent.right = y;
            }
            //更新y的父节点为x的父节点
            y.parent = x.parent;
        } else {
            this.root = y;
            this.root.parent = null;
        }

        //将y的左子节点指向x，更新x的父节点为y
        y.left = x;
        x.parent = y;
    }

    /**
     * 右旋方法
     * 右旋示意图：右旋y节点
     * <p>
     * p                       p
     * |                       |
     * y                       x
     * / \          ---->      / \
     * x   ry                  lx  y
     * / \                         / \
     * lx  ly                      ly  ry
     * <p>
     * 右旋都做了几件事？
     * 1.将x的右子节点 赋值 给了 y 的左子节点，并且更新x的右子节点的父节点为 y
     * 2.将y的父节点（不为空时）指向x，更新x的父节点为y的父节点
     * 3.将x的右子节点指向y，更新y的父节点为x
     */
    private void rightRotate(RBNode y) {
        RBNode x = y.left;
        //将x的右子节点赋值给y的左边
        y.left = x.right;
        //并且把y设置为x的右子节点的父节点
        if (x.right != null) {
            x.right.parent = y;
        }

        //将y的父节点p（非空时）指向x
        if (y.parent != null) {
            //如果y是parent左子树，则把x安放到parent的左边
            if (y.parent.left == y) {
                y.parent.left = x;
            } else {
                y.parent.right = x;
            }
            //更新x的父节点为y的父节点
            x.parent = y.parent;
        } else {
            this.root = x;
            this.root.parent = null;
        }

        //将x的右子节点赋值为y，将y的父节点设置为x
        x.right = y;
        y.parent = x;
    }

    /**
     * 获取当前节点的父节点
     */
    private RBNode parentOf(RBNode node) {
        if (node != null) {
            return node.parent;
        }
        return null;
    }

    /**
     * node节点是否为红色
     *
     * @return boolean true 表示是红色  false 表示不是红色
     */
    private boolean isRed(RBNode node) {
        if (node != null) {
            return node.isColor() == RED;
        }
        return false;
    }

    /**
     * node节点是否为黑色
     *
     * @return boolean true 表示是黑色  false 表示不是黑色
     */
    private boolean isBlack(RBNode node) {
        if (node != null) {
            return node.isColor() == BLACK;
        }
        return false;
    }

    /**
     * 设置节点为红色
     */
    private void setRed(RBNode node) {
        if (node != null) {
            node.setColor(RED);
        }
    }

    /**
     * 设置节点为黑色
     */
    private void setBlack(RBNode node) {
        if (node != null) {
            node.setColor(BLACK);
        }
    }

    /**
     * 中序打印，可以将二叉查找树有顺序的打印出来
     */
    public void inOrderPrint() {
        if (this.root != null) {
            inOrderPrint(this.root);
        }
    }

    private void inOrderPrint(RBNode node) {
        if (node != null) {
            inOrderPrint(node.left);
            System.out.println("key -> " + node.key + ", value -> " + node.value);
            inOrderPrint(node.right);
        }
    }

    /**
     * 红黑树Node
     */
    static class RBNode<K extends Comparable<K>, V> {
        //颜色
        private boolean color;
        //左子节点
        private RBNode left;
        //右子节点
        private RBNode right;
        //父节点
        private RBNode parent;
        //key
        private K key;
        //value
        private V value;

        public RBNode(boolean color, RBNode left, RBNode right, RBNode parent, K key, V value) {
            this.color = color;
            this.left = left;
            this.right = right;
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public RBNode() {
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }
}

